'''
Created on Jan 30, 2010

@author: oabalbin
'''
import os;
import numpy as np
from datetime import datetime
from collections import deque
from collections import deque, defaultdict
class record():
    
    def __init__(self, probeID, geneSym, sample, expVal):
        self.probeID = probeID
        self.geneSym = geneSym
        self.sample = sample
        self.expVal = expVal

class genrecord():
    
    def __init__(self, probeID, geneSym, sample, expVal, std):
        self.probeID = probeID
        self.geneSym = geneSym
        self.sample = sample
        self.expVal = expVal
        self.std = std


class cmprecord():
    
    def __init__(self, probeID, geneSym, sample, expVal, std):
        self.probeID = probeID
        self.geneSym = geneSym
        self.sample = sample
        self.expVal = expVal
        self.std = std


class drugBankrecord():
    #DB_card, DrugBrand_Name, CAS_Registry_Number, Generic_name, Target_geneName, Target_PDBID,Target_name
    
    def __init__(self, dbcard, brandname, casRegistry, genericname, targetGen,targetProt,targetName):
        self.dbcard = dbcard
        self.targetGen = targetGen
        self.targetProt = targetProt
        self.brandname = brandname
        self.casRegistry = casRegistry
        self.genericname = genericname
        self.targetName = targetName
    
    
class cmpName():
    
    def __init__(self,ncsid, casid, chmeName):
        
        self.ncsid=ncsid
        self.casid=casid
        #list
        self.chemName=chmeName
        
            
class drugBankrecord2():
    
    def __init__(self, dict_values):
        
        self.dbcard=dict_values['#BEGIN_DRUGCARD'][0]
        self.brandname=dict_values['Brand_Names:'][0]
        self.cas=dict_values['CAS_Registry_Number:'][0]
        self.genericname=dict_values['Generic_Name:'][0]
        self.targetgen=dict_values['Drug_Target_1_Gene_Name:'][0]
        self.targetprot=dict_values['Drug_Target_1_PDB_ID:'][0]
        self.targetprotname=dict_values['Drug_Target_1_Name:'][0]
    
    def print_myvalues(self):
        return [self.dbcard,self.brandname,self.cas,self.genericname,self.targetgen,self.targetprot,self.targetprotname]



class arrayMatrix():
    
    def __init__(self, ngenes,nsamples):
        self.probeID = deque([])
        self.geneSym = deque([]) # Note Caution: that in earliest code this was a set. Now it is a list. Then it is not a unique list of genes but the list of genes times nsamples
        self.sample = deque([])
        self.geneSynd=deque([]) #set()
        self.genInd = {}
        self.sampleInd={}
        # Arrays
        self.expValu = np.empty([ngenes,nsamples], dtype=float) 
        self.stdVal = np.empty([ngenes,nsamples], dtype=float)
        self.expmask=np.empty([ngenes,nsamples], dtype=bool)
        self.expVal = np.ma.array(self.expValu,mask=self.expmask)
    
    def create_gen_index(self):
        
        for i, g in enumerate(self.geneSynd):
            self.genInd[g[1]] = g[0]
    
    def create_sample_index(self):
        
        for sp in self.sample:
            self.sampleInd[sp[1]] = sp[0] 
    
    
    def get_thissample_indeces(self,mysamples):
        """
        It returns the specific indeces for the samples in mysamples
        """
        self.create_sample_index()
        thissamples_indices=deque()
        
        for spp in self.sample:
            sp = spp[1]
            if sp in mysamples:
                ind = self.sampleInd[sp]
                thissamples_indices.append(ind)
        
        return thissamples_indices

    
    def normalization(self):
        """
        Normalization of the data.
        It follows oncomine normalization method.
        All data sets are log-transformed and median-centered per array, 
        and standard deviations were normalized to one per array.
        """
        '''
        self.expVal = np.log2(self.expVal)
        #See comments in the use of this function
        self.addjust_expmat()                  
        # Get the mean of each array-sample. So, mean by column 
        exp_mean = np.mean(self.expVal, axis=0)
        # Get the standard deviation of each array-sample. So, mean by column
        exp_std = np.std(self.expVal, axis=0)
    
        print list(exp_mean)
        print list(exp_std)
        
        exp_norm = (self.expVal - exp_mean)/exp_std
        '''
        
        self.expVal = np.log2(self.expVal)
        #See comments in the use of this function
        self.addjust_expmat()              
        # Get the mean of each array-sample. So, mean by column 
        exp_mean = np.ma.mean(self.expVal, axis=0)
        # Get the standard deviation of each array-sample. So, mean by column
        exp_std = np.ma.std(self.expVal, axis=0)            
        exp_norm = (self.expVal - exp_mean)/exp_std

        # Normalized data.
        self.expVal = exp_norm
    
    def center_scale_norm(self):
        """
        Normalization of the data.
        It follows oncomine normalization method.
        All data sets are log-transformed and median-centered per array, 
        and standard deviations were normalized to one per array.
        
        """
                          
        # Get the mean of each array-sample. So, mean by column 
        exp_mean = np.ma.mean(self.expVal, axis=0)
        # Get the standard deviation of each array-sample. So, mean by column
        exp_std = np.ma.std(self.expVal, axis=0)            
        exp_norm = (self.expVal - exp_mean)/exp_std

        # Normalized data.
        self.expVal = exp_norm
    
    def log2_norm(self):
        """
        """
        self.expVal = np.log2(self.expVal)
        self.addjust_expmat()
        
    def quantile_norm(self):
        """
        It performs a quantile normalization of the expression data
        """
        
        self.expVal = np.log2(self.expVal)
        #See comments in the use of this function
        self.addjust_expmat()              
        
        expVal = np.ma.copy(self.expVal)
        # Sort the original array and get the mean value for each row
        expVal = np.ma.sort(expVal,axis=0)
        avg_val = np.ma.mean(expVal,axis=1)# np.ma.mean(),axis=1)
        # get the index of the sorting process
        sort_ind = np.ma.argsort(self.expVal,axis=0)
        
        # asign to each column of the explVal the average value of Xsort matrix, then 
        # re-order the matrix to original order using sort_ind
        for i in range(self.expVal.shape[1]):            
            self.expVal[:,i] = avg_val[sort_ind[:,i]]
        
    
    def normalizeexp_bynormals(self,normal_samples):
        """
        It normalizes the expression value of the array by the median expression value of the normals.
        The modified expVal matrix has the same size than the original one.
        """
        
        self.create_sample_index()
        exp_innormals=np.ma.empty([self.expVal.shape[0],len(normal_samples)])
        
        for spp in self.sample:
            sp = spp[1]
            if sp in normal_samples:
                ind = self.sampleInd[sp]
                i = normal_samples.index(sp)
                #print sp, i, ind
                #print self.expVal[:,ind]
                exp_innormals[:,i] = self.expVal[:,ind]

        median_expnormals=np.ma.median(exp_innormals,axis=1)
        median_col = np.reshape(median_expnormals,(self.expVal.shape[0],-1))
        med_new_arr = np.tile(median_col,self.expVal.shape[1])
        
        # Difference between median normal and tumors
        #print self.expVal[:,1] 
        self.expVal = self.expVal - med_new_arr
        #print median_expnormals
        #print self.expVal[:,1]
    
    def normalizeexp_bynormals_raw(self,normal_samples):
        """
        It normalizes the expression value of the array by the median expression value of the normals.
        The modified expVal matrix has the same size than the original one.
        """
        
        self.create_sample_index()
        exp_innormals=np.ma.empty([self.expVal.shape[0],len(normal_samples)])
        
        for spp in self.sample:
            sp = spp[1]
            if sp in normal_samples:
                ind = self.sampleInd[sp]
                i = normal_samples.index(sp)
                #print sp, i, ind
                #print self.expVal[:,ind]
                exp_innormals[:,i] = self.expVal[:,ind]

        median_expnormals=np.ma.median(exp_innormals,axis=1)
        median_col = np.reshape(median_expnormals,(self.expVal.shape[0],-1))
        med_new_arr = np.tile(median_col,self.expVal.shape[1])
        
        # Difference between median normal and tumors
        #print self.expVal[:,1] 
        self.expVal = np.ma.true_divide(self.expVal,med_new_arr)
        #print median_expnormals
        #print self.expVal[:,1]

    
    
    def calc_median_expression(self,other_samples):
        """
        It calculates the median expression of the expression value for each gene.
        If other_samples=True it returns a matrix of median expression and dimemsions
        self.expVal.shape[0], len(other_samples), otherwise it returns a vector.
        """
        median_expnormals=np.ma.median(self.expVal,axis=1)
        if other_samples:
            median_col = np.reshape(median_expnormals,(self.expVal.shape[0],-1))
            med_new_arr = np.tile(median_col,len(other_samples))
        else:
            med_new_arr=median_expnormals
            
        return med_new_arr
                
    
    def normalize_by_avg_expval(self):
        """
        normalize the each row by the average value of that row across all column
        """
        median_val = np.ma.median(self.expVal,axis=1)
        median_col = np.reshape(median_val,(self.expVal.shape[0],-1))
        med_new_arr = np.tile(median_col,self.expVal.shape[1])
        self.expVal = self.expVal - med_new_arr 
    
    def normalize_by_thisexpval(self, external_expression_mat):
        """
        It normalizes the expression values with an external value. 
        The two matrices should have the same dimensions. 
        e.g the median expression values in normals tissue 
        """
        self.expVal = self.expVal - external_expression_mat
        
    
    def addjust_expmat(self):
        """
        This is an special function, use only for the dataset of ERG overexpression
        in which the dyes (cye5, cye3) were swap. This function only applies to that
        samples and because the data were taken in log_old change.
        The correction was suggested by Scott Tomlins.
        """
        swapsample = ['PREC_AD_ERG1_LACZ_dyeswap', 'PREC_AD_ERG1_LACZ_dyeswap_dup', 'PREC_AD_ETV1_LACZ_dyeswap',\
                      'PREC_AD_ETV1_LACZ_dyeswap_dup','RWPE_L_ERG1_GUS_dyeswap','RWPE_L_ERG1_GUS_dyes_dup',\
                      'RWPE_L_ETV1_GUS_dyeswap','RWPE_L_ETV1_GUS_dyeswap_dup', 'CPP_RWPE19','CPP_RWPE_7','Vcap_SiRNA_ERG_nontarget_swap',
                      'Vcap_SiRNA_ERG_nontarget_swap_dup', 'LNCAP_SiRNAETV1_notarget1_swap','LNCAP_SiRNAETV1_notarget2_swap','mouse1_ERG']
        
        for sp in self.sample:
            #print sp, sp[0], self.sample
            if sp[1] in swapsample:
                #print self.expVal[:,sp[0]]
                self.expVal[:,sp[0]] = -1.0*(self.expVal[:,sp[0]]) 
                #print self.expVal[:,sp[0]]
    
    def weigthed_exp_aver(self, weigths):
        """
        It Calculates a weighted average of the gene expression in expval. 
        The weights are any float number, e.i: loads in the bfrm factor 
        Returns the weighted matrix
        weigths is a dictionary of genes, loads.
        """
        
        ind=[]
        loadvec=[]
        for gen,load in weigths.iteritems():
            try:
                ind.append(self.genInd[gen])
                loadvec.append(load)
            except KeyError:
                continue
        
        gind = np.array(ind)
        gloads= np.array(loadvec)
        # extract expVal for the genes in factor
        thmat = self.expVal[gind,:]
        # multiply expVal by the weight of each gene
        weightmat = thmat.T*gloads
        # get weighted average = pathway activity
        wavg = np.ma.mean(weightmat.T, axis=0)
        wmed = np.ma.median(weightmat.T, axis=0)
        wstd = np.ma.std(weightmat.T, axis=0)
        
        '''
        print wavg
        print wmed
        print np.max(wavg)
        print wstd 
        '''
        
        return wavg, wmed, wstd
    
    
    


class cpmMatrix():
    
    def __init__(self, ngenes,nsamples):
        self.probeID = []
        self.geneSym = []
        self.sample = []
        self.geneSynd=set()
        self.genInd = {}
        # Arrays
        self.expValu = np.empty([ngenes,nsamples], dtype=float) 
        self.stdVal = np.empty([ngenes,nsamples], dtype=float)
        self.expmask=np.empty([ngenes,nsamples], dtype=bool)
        self.expVal = np.ma.array(self.expValu,mask=self.expmask)
        
    def create_gen_index(self):
                
        for i, g in enumerate(self.geneSynd):
            self.genInd[g[1]] = g[0]
    

class bfrm_factor():
    """
    It defines a container for the bfrm factors.
    """
    def __init__(self,facnum,genedic,sdgene):
        """
        genedic has the information about the gene in the factor and the loading of that gen in that particular factor.
        sdgene has information about which genes in the factor were seed genes in the bfrm analysis
        """
        self.facnum=facnum
        self.genedic=genedic
        self.seedgene=sdgene

        



class directoryTree:
    
    def __init__(self):
        self.full_path_name=''
    
    def check_create_dir(self, full_path_name):
        if not os.path.isdir(full_path_name):
            os.mkdir( full_path_name )
    
    def make_directory_tree(self):
        
        rootdirectorypath = self.full_path_name
        
        subfolder1 =  rootdirectorypath + '/bfrm'
        subfolder2 =  rootdirectorypath + '/drugsen'
        subfolder3 =  rootdirectorypath + '/drugsen/plots'
        subfolder4 =  rootdirectorypath + '/sdv'
        '''
        subfolder5 =  rootdirectorypath + ''
        subfolder6 =  rootdirectorypath + ''
        subfolder7 =  rootdirectorypath + '' 
        subfolder8 =  rootdirectorypath + ''
        '''
        directories = [subfolder1, subfolder2, subfolder3, subfolder4] #, subfolder5, subfolder6, subfolder7, subfolder8]
        
        newdirectories = []
        for subfolder in directories:
            self.check_create_dir( subfolder )
            newdirectories.append(subfolder+'/')
            
        return newdirectories
     
    def create_tstampdir( self, full_path_name ):
        """
        creates an output folder. Name folder is the date it was created
        returns the full path to the folder and the folder name only: outfolder, t.strftime
        """
        t = datetime.now()
        outfolder = full_path_name + t.strftime("%Y_%m_%d_%H_%M")
        if not os.path.isdir(outfolder):
            os.mkdir(outfolder)
            
        self.full_path_name = outfolder
        
        return outfolder, t.strftime("%Y_%m_%d_%H_%M")

class logFile:
    
    def __init__(self,full_path_name):
        self.this = open(full_path_name,'w')
        self.retthisfile()
    
    def retthisfile(self):
        return self.this
    
    def write_list_names(self, listnames):
        
        self.this.write(",".join(map(str,listnames)).replace(',','\t')+'\n') 
    
        
class cPathway():
    """
    Class Pathway to be used in the KEGG script
    """
    def __init__(self,name, genedict):            
        self.name = name
        self.genes = genedict

class cInteraction():
    """
    Class interaction to be used in the KEGG script
    """
    def __init__(self,pathname,id1,id2,names1,names2,type_int,type_rel):   
        #item1.element_id1, gen1, item1.element_id2, gen2,type_interact,type_relation
        self.pathname = pathname
        self.item1id = id1
        self.item2id = id2
        self.item1names = names1
        self.item2names = names2
        self.type_interact = type_int
        self.type_relation = type_rel
        
    def printme_basic(self):
        relationtype = ",".join(self.type_relation).replace(',','_')
        return map(str,[self.item1id,self.item2id, self.type_interact, relationtype])
    
    def printme_plus(self):
        """
        print the nodes unfolding all info in the genes
        """
        relationtype = ",".join(self.type_relation).replace(',','_')
      
        myinfo=[]
        for gen1 in self.item1names:
            for gen2 in self.item2names:       
                myinfo.append(map(str,[self.item1id, self.item2id, gen1, gen2, self.type_interact, relationtype]))

        return myinfo 
                
                
                
                
                
                
                
                
                
                
                
        
    
    
    